from PIL import Image
import numpy as np


def generate_normal_map_for_atlas(image_path, output_path, tile_size=16, strength=1.0):
    """生成图集的法线贴图"""
    # 打开图像并转换为灰度图
    atlas = Image.open(image_path).convert('RGB')
    height_map = atlas.convert('L')
    width, height = atlas.size  # 应该是 256x256

    # 转换为 NumPy 数组
    height_array = np.array(height_map, dtype=np.float32)

    # 创建空白法线贴图数组
    normal_map = np.zeros((height, width, 3), dtype=np.uint8)

    # 遍历每个图块
    for ty in range(0, height, tile_size):
        for tx in range(0, width, tile_size):
            # 遍历图块内每个像素
            for y in range(ty, ty + tile_size):
                for x in range(tx, tx + tile_size):
                    # 图块内部坐标
                    local_y = y - ty
                    local_x = x - tx

                    center = height_array[y, x]

                    # 只在图块范围内采样
                    left = height_array[y, x - 1] if local_x > 0 else center
                    right = height_array[y, x + 1] if local_x < tile_size - 1 else center
                    top = height_array[y - 1, x] if local_y > 0 else center
                    bottom = height_array[y + 1, x] if local_y < tile_size - 1 else center

                    dx = (right - left) * strength
                    dy = (bottom - top) * strength

                    normal = np.array([-dx, -dy, 1.0])
                    normal /= np.linalg.norm(normal)

                    rgb = ((normal * 0.5 + 0.5) * 255).astype(np.uint8)
                    normal_map[y, x] = [rgb[0], rgb[1], rgb[2]]

    # 保存结果
    normal_image = Image.fromarray(normal_map)
    normal_image.save(output_path)


if __name__ == '__main__':
    # 使用示例
    generate_normal_map_for_atlas(
        '../Resources/Blocks/blocks2.png',
        'Resources/Blocks/blocks_normal.png',
        tile_size=16,
        strength=0.04
    )
